package mongodb

import (
	"context"
	"fmt"
	"time"

	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"

	commonLogger "yanxue_ai_go/pkg/common/logger"
)

// Config MongoDB配置
type Config struct {
	Host     string
	Port     int
	User     string
	Password string
	Database string
}

var (
	Client   *mongo.Client
	Database *mongo.Database
)

// Init 初始化MongoDB连接
func Init(config Config) error {
	uri := fmt.Sprintf("mongodb://%s:%s@%s:%d",
		config.User, config.Password, config.Host, config.Port)

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	client, err := mongo.Connect(ctx, options.Client().ApplyURI(uri))
	if err != nil {
		return fmt.Errorf("failed to connect to MongoDB: %w", err)
	}

	// 测试连接
	if err := client.Ping(ctx, nil); err != nil {
		return fmt.Errorf("failed to ping MongoDB: %w", err)
	}

	Client = client
	Database = client.Database(config.Database)

	commonLogger.Logger.Info("MongoDB connected successfully")
	return nil
}

// Close 关闭MongoDB连接
func Close() error {
	if Client != nil {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		return Client.Disconnect(ctx)
	}
	return nil
}

// GetClient 获取MongoDB客户端
func GetClient() *mongo.Client {
	return Client
}

// GetDatabase 获取数据库实例
func GetDatabase() *mongo.Database {
	return Database
}

// GetCollection 获取集合
func GetCollection(name string) *mongo.Collection {
	return Database.Collection(name)
}

// HealthCheck 健康检查
func HealthCheck() error {
	if Client == nil {
		return fmt.Errorf("MongoDB client not initialized")
	}

	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()

	return Client.Ping(ctx, nil)
}
